

package com.sdhs.common.utils.tree;

import cn.hutool.core.util.ObjectUtil;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 树形结构工具类，如：菜单、部门等
 *
 * @author
 * @since 1.0.0
 */
@SuppressWarnings("unchecked")
public class TreeUtils {

	/**
	 * 根据pid，构建树节点
	 */

	public static <T extends TreeNode> List<T> build(List<T> treeNodes, Object pid) {
		//pid不能为空
		// AssertUtils.isNull(pid, "pid");

		List<T> treeList = new ArrayList<>();
		for (T treeNode : treeNodes) {
			if (ObjectUtil.isNull(pid) && ObjectUtil.isNull(treeNode.getPid())) {
				treeList.add(findChildren(treeNodes, treeNode));
			} else {
				if (pid.equals(treeNode.getPid()) || String.valueOf(pid).equals(String.valueOf(treeNode.getPid()))) {
					treeList.add(findChildren(treeNodes, treeNode));
				}
			}
		}
		return treeList;
	}

	/**
	 * 查找子节点
	 */
	private static <T extends TreeNode> T findChildren(List<T> treeNodes, T rootNode) {
		for (T treeNode : treeNodes) {
			if (rootNode.getId().equals(treeNode.getPid())) {
				T children = findChildren(treeNodes, treeNode);
				if (children != null && rootNode.getChildren() == null) {rootNode.setChildren(new ArrayList<T>());}
				rootNode.getChildren().add(children);
			}
		}
		return rootNode;
	}

	/**
	 * 构建树节点
	 */
	public static <T extends TreeNode> List<T> build(List<T> treeNodes) {
		List<T> result = new ArrayList<>();

		//list转map
		Map<Object, T> nodeMap = new LinkedHashMap<>(treeNodes.size());
		for (T treeNode : treeNodes) {
			nodeMap.put(treeNode.getId(), treeNode);
		}

		for (T node : nodeMap.values()) {
			T parent = nodeMap.get(node.getPid());
			if (parent != null && !(node.getId().equals(parent.getId()))) {
				if (null != parent.getChildren()) {
					parent.getChildren().add(node);
				} else {
					List<T> children = new ArrayList<T>();
					children.add(node);
					parent.setChildren(children);
				}
			} else {
				result.add(node);
			}

		}

		return result;
	}

	/**
	 * 查找子节点
	 */
	public static <T extends TreeNode> List<T> findAllChildren(List<T> treeNodes, Object pid) {
		List<T> treeList = new ArrayList<>();
		for (T treeNode : treeNodes) {
			if (pid.equals(treeNode.getPid()) || String.valueOf(pid).equals(String.valueOf(treeNode.getPid()))) {
				treeList.add(treeNode);
				treeList.addAll(findAllChildren(treeNodes, treeNode.getId()));
			}
		}
		return treeList;
	}

	/**
	 * 查找父节点
	 */
	public static <T extends TreeNode> List<T> findAllParent(List<T> treeNodes, Object pid) {
		List<T> treeList = new ArrayList<>();
		for (T treeNode : treeNodes) {
			if (pid.equals(treeNode.getId())) {
				List<T> childList = findAllParent(treeNodes, treeNode.getPid());
				treeList.addAll(childList);
			}
		}
		return treeList;
	}

}